//
//  PacketSniffer.m
//  PacketSniffer
//
//  Created by nark on 15/04/10.
//  Copyright 2010 __MyCompanyName__. All rights reserved.
//

#import "PSPacketSniffer.h"

// global vars
static pcap_t *handle;
static PSError *currentError; 

// notifications
NSString *const PSReceiveEthernetPacketNotification = @"PSReceiveEthernetPacket";
NSString *const PSReceiveIPPacketNotification		= @"PSReceiveIPPacket";
NSString *const PSReceiveARPPacketNotification		= @"PSReceiveARPPacket";



/** 
 * This interface hides PCAP low level methods
 * used by this wrapper class like the notification
 * & delegate methods.
 */
@interface PSPacketSniffer (hidden)

/* Objective-C hidden methods*/
- (BOOL)startCaptureHandler; // the capture loop NSThread handler
- (PSPacket *)handleIP:(u_char *)_args header:(const struct pcap_pkthdr *)_pkthdr packet:(const u_char *)_packet;
- (PSEthernetPacket *)handleEthernet:(u_char *)_args header:(const struct pcap_pkthdr *)_pkthdr packet:(const u_char *)_packet;

// notification & delegate calls
- (void)notifyEthernetPacket:(PSEthernetPacket *)_packet;
- (void)notifyIPPacket:(PSIPPacket *)_packet;
- (void)notifyARPPacket:(PSARPPacket *)_packet;

/* C Low-Level functions */
void got_packet(u_char *args, const struct pcap_pkthdr *pkthdr, const u_char *packet); // pcap_loop callback function

@end


#pragma mark -
#pragma mark PacketSniffer frameworks methods

@implementation PSPacketSniffer

@synthesize device;
@synthesize filter;
@synthesize snaplen;
@synthesize promiscuous;
@synthesize timeout;
@synthesize delegate;
@synthesize isRunning;
@synthesize captureType;



- (id)initWithDevice:(PSDevice *)_device 
		 captureType:(PSCaptureType)_captureType 
	 promiscuousMode:(BOOL)_promiscuous 
		 withTimeout:(int)_timeout 
			delegate:(id<SnifferListener>)_delegate {
	
	self = [super init];
	if (self != nil) 
	{
		device			= _device;
		snaplen			= 1024;
		promiscuous		= _promiscuous;
		timeout			= _timeout;
		delegate		= _delegate;
		captureType			= _captureType;
		refToSelf		= self;
		filter			= nil;
		currentError	= nil;
		handle			= NULL;
		
		nc = [NSNotificationCenter defaultCenter];
	}
	return self;
}



- (id)initWithDevice:(PSDevice *)_device 
		 captureType:(PSCaptureType)_captureType 
	 promiscuousMode:(BOOL)_promiscuous 
		 withTimeout:(int)_timeout {
	
	return [self initWithDevice:_device
					captureType:_captureType
				promiscuousMode:_promiscuous
					withTimeout:_timeout
					   delegate:nil];
	
}

- (void) dealloc
{
	[device release];
	[refToSelf release];
	
	[super dealloc];
}




- (void)setFilter:(PSFilter *)_filter {
	
	if(filter != nil) {
		[filter release];
		filter = nil;
	}
	
	filter = _filter;
}

- (pcap_t *)handle {
	return handle;
}




- (void)startCapture {
	isRunning = YES;
	
	[NSThread detachNewThreadSelector:@selector(startCaptureHandler)
							 toTarget:self
						   withObject:nil];
}


- (BOOL)stopCapture {
	
	pcap_breakloop(handle);
	
	return YES;
}

@end




#pragma mark -
#pragma mark PacketSniffer hidden methods

@implementation PSPacketSniffer (hidden)

- (BOOL)startCaptureHandler {
	NSAutoreleasePool *pool = [[NSAutoreleasePool alloc] init];
	
	NS_DURING
	
	/* PCAP open live  */
	handle = pcap_open_live([[device name] UTF8String],
							BUFSIZ, 
							((promiscuous) ? 1 : 0), 
							timeout, 
							errbuf);
	
	/* Check if handle is open live on device */
	if(handle == NULL) {
		NSString *message = [NSString stringWithFormat:@"Couldn't open device %s: %s", [[device name] UTF8String], errbuf];
		PSError *error = [[PSError alloc] initWithType:PSErrorTypeCritical
											   message:message 
												 infos:nil];
		
		[PSError notifyError:error];
		return FALSE;
	}
	
	
	/* Set PCAP non-blocking mode */
	//	if(pcap_setnonblock(handle, 0, errbuf) == -1) {
	//		fprintf(stderr, "Couldn't set non block capture.");
	//        return FALSE;
	//	}
	
	
	/* If filter, apply it */
	if(filter != nil) {
		struct bpf_program fp;
		
		/* Compile the filter */
		if (pcap_compile(handle, &fp, [[filter filter] UTF8String], 0, [device net]) == -1) {
			NSString *message = [NSString stringWithFormat:@"Couldn't parse filter %s: %s", [[filter filter] UTF8String], pcap_geterr(handle)];
			PSError *error = [[PSError alloc] initWithType:PSErrorTypeCritical
												   message:message 
													 infos:nil];
			[PSError notifyError:error];
			return FALSE;
		}
		
		if (pcap_setfilter(handle, &fp) == -1) {
			NSString *message = [NSString stringWithFormat:@"Couldn't install filter %s: %s", [[filter filter] UTF8String], pcap_geterr(handle)];
			PSError *error = [[PSError alloc] initWithType:PSErrorTypeCritical
												   message:message 
													 infos:nil];
			[PSError notifyError:error];
			return FALSE;
		}
	}
	
	/* Run capture packets loop.
	 * The following in the next method which
	 * is the PCAP loop callback function.
	 */
    if(pcap_loop(handle, -1, got_packet, NULL) == -1) {
		NSString *message = [NSString stringWithFormat:@"Couldn't not capture packets."];
		PSError *error = [[PSError alloc] initWithType:PSErrorTypeFatal
											   message:message 
												 infos:nil];
		[PSError notifyError:error];
        return FALSE;
    }
	
	
	NS_HANDLER
	NS_ENDHANDLER
	
	[pool release];
	return TRUE;
}


/** 
 * The PCAP loop callback function.
 * This is a low level C function that use the
 * ugly refToSelf var decause of broken scope.
 * It switches around the desired capture type
 * to notify delegate object about new received
 * packets.
 */
void got_packet(u_char *args, const struct pcap_pkthdr *pkthdr, const u_char *packet) {
	
	NSAutoreleasePool *pool = [[NSAutoreleasePool alloc] init];
	PSEthernetPacket *ethernetPacket = nil;
	
	NS_DURING
	
	if(refToSelf) {
		// handle the ethernet packet 
		ethernetPacket = [refToSelf handleEthernet:args header:pkthdr packet:packet];
		
		if(ethernetPacket) {
			// check the desired capture type
			switch ([refToSelf captureType]) {
					
					
				case PSEveryCaptureType: {
					// notify ip packets
					if([ethernetPacket type] == ETHERTYPE_IP) {
						PSIPPacket *ipPacket = [[PSIPPacket alloc] initWithLength:pkthdr->caplen];
						ipPacket = (PSIPPacket *)[refToSelf handleIP:args header:pkthdr packet:packet];
						[ipPacket setParent:ethernetPacket]; // set the ethernet parent frame
						[[ethernetPacket children] addObject:ipPacket]; // set the ip child frame
						[refToSelf notifyIPPacket:ipPacket];
					}
					// notify arp packets
					if(([ethernetPacket type] == ETHERTYPE_ARP) || 
					   ([ethernetPacket type] == ETHERTYPE_REVARP)) {
						PSARPPacket *arpPacket = [[PSARPPacket alloc] initWithLength:pkthdr->caplen];
						[arpPacket setParent:ethernetPacket]; // set the ethernet parent frame
						[[ethernetPacket children] addObject:arpPacket]; // set the arp child frame
						[refToSelf notifyARPPacket:arpPacket];
					}
					
					// notify ethernet packets (no control, always exist)
					[refToSelf notifyEthernetPacket:ethernetPacket];
				} break;
					
					
					
				case PSEthernetOnlyCaptureType: {
					[refToSelf notifyEthernetPacket:ethernetPacket];
				} break;
					
					
					
				case PSIPOnlyCaptureType: {
					// notify ip packets
					if([ethernetPacket type] == ETHERTYPE_IP) {
						PSIPPacket *ipPacket = [[PSIPPacket alloc] initWithLength:pkthdr->caplen];
						ipPacket = (PSIPPacket *)[refToSelf handleIP:args header:pkthdr packet:packet];
						[ipPacket setParent:ethernetPacket]; // set the ethernet parent frame
						[[ethernetPacket children] addObject:ipPacket]; // set the ip child frame
						[refToSelf notifyIPPacket:ipPacket];
					}
				} break;
					
					
					
				case PSARPOnlyCaptureType: {
					if(([ethernetPacket type] == ETHERTYPE_ARP) || 
					   ([ethernetPacket type] == ETHERTYPE_REVARP)) {
						PSARPPacket *arpPacket = [[PSARPPacket alloc] initWithLength:pkthdr->caplen];
						[arpPacket setParent:ethernetPacket]; // set the ethernet parent frame
						[[ethernetPacket children] addObject:arpPacket]; // set the arp child frame
						[refToSelf notifyARPPacket:arpPacket];
					}
				} break;
			}
		}
	}
	
	NS_HANDLER
	NS_ENDHANDLER
	
	[pool release];
}



- (void)notifyEthernetPacket:(PSEthernetPacket *)_packet {
	// post notification the the default center
	NSMutableDictionary *userInfo = [[NSMutableDictionary alloc] init];
	[userInfo setValue:_packet forKey:@"packet"];
	[nc postNotificationName:PSReceiveEthernetPacketNotification object:self userInfo:userInfo];
	
	// if delegate call it
	if(delegate) {
		if([(id)delegate respondsToSelector:@selector(receivePacket:)]) 
			[delegate receivePacket:_packet];
	}
}


- (void)notifyIPPacket:(PSIPPacket *)_packet {
	// post notification the the default center
	NSMutableDictionary *userInfo = [[NSMutableDictionary alloc] init];
	[userInfo setValue:_packet forKey:@"packet"];
	[nc postNotificationName:PSReceiveIPPacketNotification object:self userInfo:userInfo];
	
	// if delegate call it
	if(delegate)
		if([(id)delegate respondsToSelector:@selector(receivePacket:)]) {
			[delegate receivePacket:_packet];
		}
}


- (void)notifyARPPacket:(PSARPPacket *)_packet {
	// post notification the the default center
	NSMutableDictionary *userInfo = [[NSMutableDictionary alloc] init];
	[userInfo setValue:_packet forKey:@"packet"];
	[nc postNotificationName:PSReceiveARPPacketNotification object:self userInfo:userInfo];
	
	// if delegate call it
	if(delegate)
		if([(id)delegate respondsToSelector:@selector(receivePacket:)]) {
			[delegate receivePacket:_packet];
		}
}


/** 
 * Handle IP packets
 */
- (PSPacket *)handleIP:(u_char *)_args header:(const struct pcap_pkthdr *)_pkthdr packet:(const u_char *)_packet {
	NSAutoreleasePool *pool = [[NSAutoreleasePool alloc] init];
	PSIPPacket *packet = [[PSIPPacket alloc] init];
	
	NS_DURING
	
	// build an IP packet object by casting PCAP data (NEED MORE SAFE CODE !!!)
	struct nread_ip *ip = (struct nread_ip*)(_packet + 14);	// The +14 is the size of the ethernet frame
	
	packet.sourceIP			= [NSString stringWithUTF8String:inet_ntoa(ip->ip_src)];
	packet.destinationIP	= [NSString stringWithUTF8String:inet_ntoa(ip->ip_dst)];
	packet.version			= IP_V(ip);
	packet.headerLength		= IP_HL(ip);
	packet.totalLength		= ip->ip_len;
	//packet.
	
	// here we need to test protocol (TCP/UDP...)
	if (ip->ip_p == IPPROTO_TCP) { // TCP
		
		PSTCPHeader *tcpHeader = [[PSTCPHeader alloc] init];
		struct nread_tcp *tcp = (struct nread_tcp*)(packet + 14 + 20);
		[[packet children] addObject:tcpHeader];
	
		// build tcp header
		tcpHeader.parent				= packet;
		tcpHeader.checksum				= [NSString stringWithFormat:@"%u", ntohs(tcp->th_sum)];
		tcpHeader.urgentPointer			= [NSString stringWithFormat:@"%u", ntohs(tcp->th_urp)];
		tcpHeader.dataOffset			= ntohs(tcp->th_off);
		tcpHeader.sequenceNumber		= ntohs(tcp->th_seq);
		tcpHeader.acknowledgementNumber = ntohs(tcp->th_ack);
		tcpHeader.window				= ntohs(tcp->th_win);
		tcpHeader.sourcePort			= [NSString stringWithFormat:@"%u", ntohs(tcp->th_sport)];
		tcpHeader.destinationPort		= [NSString stringWithFormat:@"%u", ntohs(tcp->th_dport)];
		
	} else if (ip->ip_p == IPPROTO_UDP) { // UDP
		
		PSUDPHeader *udpHeader	= [[PSUDPHeader alloc] init];
		
		int udp_size			= sizeof(struct nread_udp*);
		struct nread_udp *udp	= (struct nread_udp*)(packet + 14 + udp_size);
		[[packet children] addObject:udpHeader];
		
		udpHeader.parent			= packet;
		udpHeader.sourcePort		= udp->uh_sport;
		udpHeader.destinationPort	= udp->uh_dport;
		udpHeader.udpLength			= ntohs(udp->uh_ulen);
		udpHeader.checksum			= [NSString stringWithFormat:@"%u", ntohs(udp->uh_sum)];
		
	} else if (ip->ip_p == IPPROTO_ICMP) {
	
		
		
	} else {
		// send an unsupported ip protocol error
		NSString *message = [NSString stringWithFormat:@"Unsoported IP Protocol header."];
		PSError *error = [[PSError alloc] initWithType:PSErrorTypeWarning
											   message:message 
												 infos:nil];
		[PSError notifyError:error];
	}
	
	NS_HANDLER
	NS_ENDHANDLER
	
	[pool release];
	return packet;
}


/** 
 * Handle Ethernet packets.
 */
- (PSEthernetPacket *)handleEthernet:(u_char *)_args header:(const struct pcap_pkthdr *)_pkthdr packet:(const u_char *)_packet {
	NSAutoreleasePool *pool = [[NSAutoreleasePool alloc] init];
	PSEthernetPacket *etherPacket = [[PSEthernetPacket alloc] initWithLength:_pkthdr->len];
	
	u_short ether_type;
	char *ether_src;
	char *ether_dst;
	
	NS_DURING
	
    u_int caplen = _pkthdr->caplen;
    struct ether_header *eptr;  /* net/ethernet.h */
	
    if (caplen < ETHER_HDRLEN)
    {
		NSString *message = [NSString stringWithFormat:@"Packet length less than ethernet header length."];
		PSError *error = [[PSError alloc] initWithType:PSErrorTypeFatal
											   message:message 
												 infos:nil];
		[PSError notifyError:error];
        return nil;
    }
	
    /* build ethernet packet object by casting PCAP data  */
    eptr = (struct ether_header *)_packet;
    ether_type = ntohs(eptr->ether_type);
	ether_src  = ether_ntoa((struct ether_addr*)eptr->ether_shost);
	ether_dst  = ether_ntoa((struct ether_addr*)eptr->ether_dhost);
	
	NSLog(@"src : %s", ether_src);
	NSLog(@"dst : %s", ether_dst);
	
	etherPacket.type					= ether_type;
	etherPacket.sourceMACAddress		= [NSString stringWithUTF8String:ether_src];
	etherPacket.destinationMACAddress	= [NSString stringWithUTF8String:ether_dst];
	etherPacket.dateReceived			= [NSDate date];
	etherPacket.capturedLength			= _pkthdr->caplen;
	
	NS_HANDLER
	NS_ENDHANDLER
	
	[pool release];
	
    return etherPacket;
}


@end
